package class11;

import java.util.*;

/**
 * @author zhangchaoliang
 * create 2022
 */
public class Prim3 {

    public static class EdgeComparator implements Comparator<Edge>{

        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }
    }

    public static Set<Edge> primMst(Graph graph){
        PriorityQueue<Edge> priorityQueue =
                new PriorityQueue<>(new EdgeComparator());
        HashSet<Node> nodeSet = new HashSet<>();
        HashSet<Edge> edgeSet = new HashSet<>();

        Set<Edge> result = new HashSet<>();

        for (Node node:graph.nodes.values()){
            if (!nodeSet.contains(node)) {

                nodeSet.add(node);
                for (Edge edge:node.edges){
                    if (!edgeSet.add(edge)){
                        edgeSet.add(edge);
                        priorityQueue.add(edge);
                    }
                }

                while (!priorityQueue.isEmpty()){
                    Edge edge = priorityQueue.poll();
                    Node toNode = edge.to;
                    if (!nodeSet.contains(toNode)){
                        nodeSet.add(toNode);
                        result.add(edge);
                        for (Edge nextEdge : toNode.edges){
                            if (!edgeSet.contains(nextEdge)){
                                edgeSet.add(nextEdge);
                                priorityQueue.add(nextEdge);
                            }
                        }
                    }
                }


            }

            break;
        }
        return result;
    }
}
